{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ef331be9",
   "metadata": {
    "id": "ef331be9"
   },
   "source": [
    "![nebullvm nebuly AI accelerate inference optimize DeepLearning](https://user-images.githubusercontent.com/38586138/201391643-a80407e5-2c28-409c-90c9-327795cd27e8.png)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f260653a",
   "metadata": {
    "id": "f260653a"
   },
   "source": [
    "# Accelerate Stable Diffusion with Speedster\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bdf3af5",
   "metadata": {
    "id": "8bdf3af5"
   },
   "source": [
    "Hi and welcome 👋\n",
    "\n",
    "In this notebook we will discover how in just a few steps you can speed up the response time of deep learning model inference using the Speedster app from the open-source library nebullvm.\n",
    "\n",
    "With Speedster's latest API, you can speed up models up to 10 times without any loss of accuracy (option A), or accelerate them up to 20-30 times by setting a self-defined amount of accuracy/precision that you are willing to trade off to get even lower response time (option B). To accelerate your model, Speedster takes advantage of various optimization techniques such as deep learning compilers (in both option A and option B), quantization, half accuracy, and so on (option B).\n",
    "\n",
    "Let's jump to the code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cXXh1ifQ13mH",
   "metadata": {
    "id": "cXXh1ifQ13mH"
   },
   "source": [
    "# Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48aljCHu14-H",
   "metadata": {
    "id": "48aljCHu14-H"
   },
   "source": [
    "Install Speedster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "QFQh3BVr1-GO",
   "metadata": {
    "id": "QFQh3BVr1-GO"
   },
   "outputs": [],
   "source": [
    "!pip install speedster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a7a86b3",
   "metadata": {
    "id": "8a7a86b3"
   },
   "source": [
    "Install deep learning compilers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cffbfa32",
   "metadata": {
    "id": "cffbfa32"
   },
   "outputs": [],
   "source": [
    "!python -m nebullvm.installers.auto_installer --frameworks diffusers --compilers all"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73072506",
   "metadata": {
    "id": "73072506"
   },
   "source": [
    "## Model and Dataset setup"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aeb2c521",
   "metadata": {},
   "source": [
    "First of all we have to choose the version of Stable Diffusion we want to optimize, speedster officially supports the most used versions:\n",
    "- `CompVis/stable-diffusion-v1-4`\n",
    "- `runwayml/stable-diffusion-v1-5`\n",
    "- `stabilityai/stable-diffusion-2-1-base`\n",
    "\n",
    "Other Stable Diffusion versions from the Diffusers library should work but have never been tested. If you try a version not included among these and it works, please feel free to report it to us on [Discord](https://discord.com/invite/RbeQMu886J) so we can add it to the list of supported versions. If you try a version that does not work, you can open an issue and possibly a PR on [GitHub](https://github.com/nebuly-ai/nebullvm/issues)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e4d55115",
   "metadata": {
    "id": "e4d55115"
   },
   "source": [
    "For this notebook, we are going to select Stable Diffusion 1.4. Let's download and load it using the diffusers API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d633cf21",
   "metadata": {
    "id": "d633cf21",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import StableDiffusionPipeline\n",
    "\n",
    "# Select Stable Diffusion version\n",
    "model_id = \"CompVis/stable-diffusion-v1-4\"\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "if device == \"cuda\":\n",
    "    # On GPU we load by default the model in half precision, because it's faster and lighter.\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)\n",
    "else:\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f7653250",
   "metadata": {},
   "source": [
    "We can easily test the loaded model by generating a sample image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b640885",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_prompt = \"futuristic llama with a cyberpunk city on the background\"\n",
    "\n",
    "pipe.to(device)\n",
    "pipe(test_prompt).images[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "11aa0739",
   "metadata": {
    "id": "11aa0739"
   },
   "source": [
    "Let's now create an example dataset with some random sentences, that will be used later for the optimization process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbbfeeb2",
   "metadata": {
    "id": "cbbfeeb2"
   },
   "outputs": [],
   "source": [
    "input_data = [\n",
    "    \"a photo of an astronaut riding a horse on mars\",\n",
    "    \"a monkey eating a banana in a forest\",\n",
    "    \"white car on a road surrounded by palm trees\",\n",
    "    \"a fridge full of bottles of beer\",\n",
    "    \"madara uchiha throwing asteroids against people\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17040431",
   "metadata": {
    "id": "17040431"
   },
   "source": [
    "## Speed up inference with Speedster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44ddc21d",
   "metadata": {
    "id": "44ddc21d"
   },
   "source": [
    "It's now time of improving a bit the performance in terms of speed. Let's use `Speedster`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9d934f6",
   "metadata": {
    "id": "f9d934f6"
   },
   "outputs": [],
   "source": [
    "from speedster import optimize_model, save_model, load_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76248033",
   "metadata": {
    "id": "76248033"
   },
   "source": [
    "Using Speedster is very simple and straightforward! Just use the `optimize_model` function and provide as input the model, some input data as example and the optimization time mode. Optionally a dynamic_info dictionary can be also provided, in order to support inputs with dynamic shape."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "75b339c3",
   "metadata": {},
   "source": [
    "**Optimisation of stable diffusion requires a lot of RAM. If you are running this notebook on google colab, make sure to use the high RAM option, otherwise the kernel may crash. If the kernel crashes also when using the high RAM option, please try adding also `\"torchscript\"` to the `ignore_compilers` list.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "zPC_EDwEJIM0",
   "metadata": {
    "id": "zPC_EDwEJIM0"
   },
   "outputs": [],
   "source": [
    "optimized_model = optimize_model(\n",
    "    model=pipe,\n",
    "    input_data=input_data,\n",
    "    optimization_time=\"unconstrained\",\n",
    "    ignore_compilers=[\"torch_tensor_rt\", \"tvm\"],  # Some compilers have issues with Stable Diffusion, so it's better to skip them.\n",
    "    metric_drop_ths=0.1,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fdae59d2",
   "metadata": {},
   "source": [
    "If running on GPU, here you should obtain a speedup of about 60% on the UNet. We run the optimization on a **3090Ti** and here are our results:\n",
    "- **Original Model (PyTorch, fp16): 51,298 ms/batch**\n",
    "- **Optimized Model (TensorRT, fp16): 32,164 ms/batch**\n",
    "\n",
    "If the optimized model you obtained is not a TensorRT one, probably there was an error during the optimization. If running on colab, it could happen that the standard gpu is not enough to run the optimization, so we suggest to select a premium gpu with more memory.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "af9f86ac",
   "metadata": {},
   "source": [
    "If everything worked correctly, let's check the output of the optimized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa443637",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model(test_prompt).images[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6e5b3b21",
   "metadata": {
    "id": "6e5b3b21"
   },
   "source": [
    "Let's run the prediction 10 times to calculate the average response time of the original model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09170c78",
   "metadata": {},
   "outputs": [],
   "source": [
    "if device == \"cuda\":\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)\n",
    "else:\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id)\n",
    "\n",
    "pipe.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3bc5c98",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d3bc5c98",
    "outputId": "e0596cf2-fa96-4c50-c012-f5cdab82e681"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "times = []\n",
    "\n",
    "# Warmup for 2 iterations\n",
    "for _ in range(2):\n",
    "    with torch.no_grad():\n",
    "        final_out = pipe(test_prompt).images[0]\n",
    "\n",
    "# Benchmark\n",
    "for _ in range(8):\n",
    "    st = time.time()\n",
    "    with torch.no_grad():\n",
    "        final_out = pipe(test_prompt).images[0]\n",
    "    times.append(time.time()-st)\n",
    "original_model_time = sum(times)/len(times)\n",
    "print(f\"Average response time for original Stable Diffusion 1.4: {original_model_time} s\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3db0a7a1",
   "metadata": {
    "id": "3db0a7a1"
   },
   "source": [
    "Let's run the prediction 10 times to calculate the average response time of the optimized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3e83997",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "a3e83997",
    "outputId": "7a416b14-f170-4df9-d416-026f06a7d980"
   },
   "outputs": [],
   "source": [
    "times = []\n",
    "\n",
    "for _ in range(2):\n",
    "    with torch.no_grad():\n",
    "        final_out = optimized_model(test_prompt).images[0]\n",
    "\n",
    "# Benchmark\n",
    "for _ in range(8):\n",
    "    st = time.time()\n",
    "    with torch.no_grad():\n",
    "        final_out = optimized_model(test_prompt).images[0]\n",
    "    times.append(time.time()-st)\n",
    "optimized_model_time = sum(times)/len(times)\n",
    "print(f\"Average response time for optimized Stable Diffusion 1.4: {optimized_model_time} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb60d8c",
   "metadata": {
    "id": "ceb60d8c"
   },
   "source": [
    "## Save and reload the optimized model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9eda1a0",
   "metadata": {},
   "source": [
    "We can easily save to disk the optimized model with the following line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b6fcbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_model(optimized_model, \"model_save_path\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c968d51",
   "metadata": {},
   "source": [
    "We can then load again the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1340c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model = load_model(\"model_save_path\", pipe=pipe)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9d94ded9",
   "metadata": {},
   "source": [
    "## Advanced: Further increase performance with TensorRT Plugins (GPU Only)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bcd2bd77",
   "metadata": {},
   "source": [
    "Reference: https://github.com/NVIDIA/TensorRT/tree/main/demo/Diffusion"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4f9e28ee",
   "metadata": {},
   "source": [
    "To achieve the best results on GPU, we have to activate the TensorRT Plugins. `Speedster` supports their usage on Stable Diffusion models from version `0.9.0`.\n",
    "\n",
    "First of all, install them by following this [guide](https://github.com/nebuly-ai/nebullvm/tree/main/notebooks/speedster/diffusers#setup-tensorrt-plugins-optional), then edit the cell below with the correct paths.\n",
    "\n",
    "If you are working in the nebullvm docker image, the plugins are installed and activated by default, so the optimization results of this section should be the same as in the previous one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf673d58",
   "metadata": {},
   "outputs": [],
   "source": [
    "plugin_path = \"/content/TensorRT/build/out/libnvinfer_plugin.so\"  # EDIT THIS PATH\n",
    "tensorrt_lib_path = \"/content/TensorRT-8.5.3.1/lib\"  # EDIT THIS PATH\n",
    "\n",
    "import os\n",
    "if not os.path.exists(plugin_path) or not os.path.exists(tensorrt_lib_path):\n",
    "    raise Exception(\"The paths provided above are invalid, please edit them according to your configuration.\")\n",
    "\n",
    "os.environ[\"LD_PRELOAD\"] = plugin_path\n",
    "if \"LD_LIBRARY_PATH\" not in os.environ:\n",
    "    os.environ[\"LD_LIBRARY_PATH\"] = \"\"\n",
    "os.environ[\"LD_LIBRARY_PATH\"] = os.environ[\"LD_LIBRARY_PATH\"] + \":\" + tensorrt_lib_path\n",
    "os.environ[\"CUDA_MODULE_LOADING\"] = \"LAZY\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a3103f7a",
   "metadata": {},
   "source": [
    "Let's repeat again the optimization process, this time we will use only TensorRT from the ONNX pipeline for faster optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140a86d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available() is False:\n",
    "    raise Exception(\"You are running in a CPU-only machine, TensorRT can be used only on GPU.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb603f8c",
   "metadata": {},
   "source": [
    "**Optimisation of stable diffusion requires a lot of RAM. If you are running this notebook on google colab, make sure you use the high RAM option, otherwise the kernel may crash.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32c0f2ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model = optimize_model(\n",
    "    model=pipe,\n",
    "    input_data=input_data,\n",
    "    optimization_time=\"unconstrained\",\n",
    "    ignore_compilers=[\"torch_tensor_rt\", \"onnxruntime\", \"torchscript\", \"tvm\"],\n",
    "    metric_drop_ths=0.1,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dc247377",
   "metadata": {},
   "source": [
    "If running on GPU, here you should obtain a speedup of about 160% on the UNet. We run the optimization on a **3090Ti** and here are our results:\n",
    "- **Original Model (PyTorch, fp16): 51,298 ms/batch**\n",
    "- **Optimized Model (TensorRT with Plugins, fp16): 19,8 ms/batch**\n",
    "\n",
    "If again the optimized model you obtained is not a TensorRT one, probably there was an error during the optimization. If running on colab, it could happen that the standard gpu is not enough to run the optimization, so we suggest to select a premium gpu with more memory.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b0921cc6",
   "metadata": {},
   "source": [
    "Let's run the prediction 10 times to calculate the average response time of the optimized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6bb526",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "times = []\n",
    "\n",
    "for _ in range(2):\n",
    "    with torch.no_grad():\n",
    "        final_out = optimized_model(test_prompt).images[0]\n",
    "\n",
    "# Benchmark\n",
    "for _ in range(8):\n",
    "    st = time.time()\n",
    "    with torch.no_grad():\n",
    "        final_out = optimized_model(test_prompt).images[0]\n",
    "    times.append(time.time()-st)\n",
    "optimized_model_time = sum(times)/len(times)\n",
    "print(f\"Average response time for optimized Stable Diffusion 1.4: {optimized_model_time} s\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5021274b",
   "metadata": {},
   "source": [
    "Finally, we can use the optimized model to generate a sample image and see the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d20e5fce",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model(test_prompt).images[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb234e5e",
   "metadata": {
    "id": "cb234e5e"
   },
   "source": [
    "Great! Was it easy? How are the results? Do you have any comments?\n",
    "Share your optimization results and thoughts with <a href=\"https://discord.gg/RbeQMu886J\" target=\"_blank\"> our community on Discord</a>, where we chat about Speedster and AI acceleration.\n",
    "\n",
    "Note that the acceleration of Speedster depends very much on the hardware configuration and your AI model. Given the same input model, Speedster can accelerate it by 10 times on some machines and perform poorly on others.\n",
    "\n",
    "If you want to learn more about how Speedster works, look at other tutorials and performance benchmarks, check out the links below or write to us on Discord."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b77ff2ac",
   "metadata": {
    "id": "b77ff2ac"
   },
   "source": [
    "<center> \n",
    "    <a href=\"https://discord.com/invite/RbeQMu886J\" target=\"_blank\" style=\"text-decoration: none;\"> Join the community </a> |\n",
    "    <a href=\"https://nebuly.gitbook.io/nebuly/welcome/questions-and-contributions\" target=\"_blank\" style=\"text-decoration: none;\"> Contribute to the library </a>\n",
    "</center>\n",
    "\n",
    "<center> \n",
    "    <a href=\"https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/speedster#key-concepts\" target=\"_blank\" style=\"text-decoration: none;\"> How speedster works </a> •\n",
    "    <a href=\"https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/speedster#documentation\" target=\"_blank\" style=\"text-decoration: none;\"> Documentation </a> •\n",
    "    <a href=\"https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/speedster#quick-start\" target=\"_blank\" style=\"text-decoration: none;\"> Quick start </a> \n",
    "</center>"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "4ca44071b2152bc556aa4c839392f76fd4b80aa39d34257f2d304fa0d1d8b7d9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
